# -*- coding: utf-8 -*-
"""
Created on Sun Feb 21 21:52:33 2021

@author: zhiyinan

RNN Data Generation - 1D irrigation
"""

import numpy as np 
from casadi import *
from Parameters1D_casadi import *
import van_Genuchten1D_casadi as vg 
import matplotlib.pyplot as plt
from Model1D_casadi_data import *
import mpctools as mpc
import time

soilPars = loamySoil()#The soil parameters
totalDepth, axialNodes, axialDistance, totalNodes = spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = temporalVariable() #time parameters

irr_freq = 48*60*60/samplingTime
lbu = -5.10e-05  #  unit [m/s]
ubu = -4.94e-05  #  unit [m/s]
   
# lbu = -4.9e-05  #  unit [m/s]
# ubu = -4.7e-05  #  unit [m/s]
u_value = (ubu - lbu) * np.random.random_sample((int(totTimeSteps/irr_freq), 1)) + lbu
u_train  = np.zeros((totTimeSteps,1))
for i in range(int(totTimeSteps/irr_freq)):
    u_train[i*int(irr_freq)] = u_value[i]
# np.savetxt('u_RNN_train.txt', u_train)

# initTime = time.process_time()
# u_train = np.loadtxt('u_RNN_train.txt')


initCondition = 1.2*-0.8*np.ones(totalNodes)

#Arrays to store the results
headArray = np.zeros((totTimeSteps+1, totalNodes)) # Pressure head
headArray[0,:] = initCondition

volMoisture = np.zeros((totTimeSteps+1, totalNodes)) #Volumetric Mositure
volMoisture[0,:] = volMoistureAllNodes_1D(initCondition).full().ravel()


##Creating the Casadi integrator
#Symbolic variables

# headSym = SX.sym('h', totalNodes)
# inpSym = SX.sym('u', 1)

#Symbolic polar model
# polarModelSym = RichardsPolar_1D(headSym, inpSym, soilPars)

ode_sim = mpc.DiscreteSimulator(RichardsPolar_1D, samplingTime, [totalNodes, 1], ['x', 'u'])
# ODE = {'x': headSym, 'p': vertcat(inpSym), 'ode': polarModelSym}
# opts = {'tf': samplingTimeInternal, 'regularity_check': True}

# I = integrator('I', 'cvodes', ODE, opts)


# def smallerTimeSteps(initHead, irrigAmnt):
#     interHead = initHead
#     for i in range(internalTimeSteps):
#         intRes = I(x0=interHead, p=vertcat(irrigAmnt))
#         interHead = intRes['xf'].full().ravel()
#     return interHead


for i in range(1, 100): #totTimeSteps + 1


    #intRes = I(x0=headArray[i-1,:], p=vertcat(u[i-1], procNoise[i-1], cropCoeff[i-1], refEvap[i-1]))
    #headArray[i,:] = intRes['xf'].full().ravel()
    headArray[i,:] = ode_sim.sim(headArray[i-1,:], u_train[i-1])
    volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:]).full().ravel()

